#include "RasterizeGaussians.h"
#include "GSSplat.h"

//******************************************************************
//FUNCTION:
std::tuple<torch::Tensor,
    torch::Tensor,
    torch::Tensor,
    torch::Tensor,
    torch::Tensor> binAndSortGaussians(int vNumPoints, int vNumIntersects,
        torch::Tensor vXYs,
        torch::Tensor vDepths,
        torch::Tensor vRadii,
        torch::Tensor vCumTilesHit,
        TileBounds vTileBounds) {
    auto ToIntersectGaussian = map_gaussian_to_intersects_tensor(vNumPoints, vNumIntersects,
        vXYs, vDepths, vRadii, vCumTilesHit, vTileBounds);

    //Authorm comment:  unique IDs for each gaussian in the form (tile | depth id)
    torch::Tensor IsectIds = std::get<0>(ToIntersectGaussian);

    //Authorm comment:  Tensor that maps isect_ids back to cumHitTiles
    torch::Tensor GaussianIds = std::get<1>(ToIntersectGaussian);

    auto Sorted = torch::sort(IsectIds);

    //Authorm comment:  sorted unique IDs for each gaussian in the form (tile | depth id)
    torch::Tensor IsectIdsSorted = std::get<0>(Sorted);
    torch::Tensor SortedIndices = std::get<1>(Sorted);

    //Authorm comment:  sorted Tensor that maps isect_ids back to cumHitTiles
    torch::Tensor GaussianIdsSorted = torch::gather(GaussianIds, 0, SortedIndices);

    //Authorm comment:  range of gaussians hit per tile
    torch::Tensor TileBins = get_tile_bin_edges_tensor(vNumIntersects, IsectIdsSorted);
    return std::make_tuple(IsectIds, GaussianIds, IsectIdsSorted, GaussianIdsSorted, TileBins);
}

//******************************************************************
//FUNCTION:
torch::Tensor CRasterizeGaussians::forward(AutogradContext* vCTx,
    torch::Tensor vXYs,
    torch::Tensor vDepths,
    torch::Tensor vRadii,
    torch::Tensor vConics,
    torch::Tensor vNumTilesHit,
    torch::Tensor vColors,
    torch::Tensor vOpacity,
    int vImgHeight,
    int vImgWidth,
    torch::Tensor vBackground
) {

    int NumPoints = vXYs.size(0);

    TileBounds TileBounds = std::make_tuple(
        (vImgWidth + BLOCK_X - 1) / BLOCK_X,
        (vImgHeight + BLOCK_Y - 1) / BLOCK_Y,
        1
    );
    std::tuple<int, int, int> Block = std::make_tuple(BLOCK_X, BLOCK_Y, 1);
    std::tuple<int, int, int> ImgSize = std::make_tuple(vImgWidth, vImgHeight, 1);

    torch::Tensor CumTilesHit = torch::cumsum(vNumTilesHit, 0, torch::kInt32);
    int NumIntersects = CumTilesHit[CumTilesHit.size(0) - 1].item<int>();

    auto Bin = binAndSortGaussians(NumPoints, NumIntersects, vXYs, vDepths, vRadii, CumTilesHit, TileBounds);
    torch::Tensor GaussianIdsSorted = std::get<3>(Bin);
    torch::Tensor TileBins = std::get<4>(Bin);

    auto t = rasterize_forward_tensor(TileBounds, Block, ImgSize,
        GaussianIdsSorted,
        TileBins,
        vXYs,
        vConics,
        vColors,
        vOpacity,
        vBackground);
    //Authorm comment:  Final image
    torch::Tensor OutImg = std::get<0>(t);

    torch::Tensor FinalTs = std::get<1>(t);

    //Authorm comment:  Map of tile bin IDs
    torch::Tensor FinalIdx = std::get<2>(t);

    vCTx->saved_data["imgWidth"] = vImgWidth;
    vCTx->saved_data["imgHeight"] = vImgHeight;
    vCTx->save_for_backward({ GaussianIdsSorted, TileBins, vXYs, vConics, vColors, vOpacity, vBackground, FinalTs, FinalIdx });

    return OutImg;
}

//******************************************************************
//FUNCTION:
tensor_list CRasterizeGaussians::backward(AutogradContext* vCTx, tensor_list vGradOutputs) {
    torch::Tensor OutImg = vGradOutputs[0];
    int ImgHeight = vCTx->saved_data["imgHeight"].toInt();
    int ImgWidth = vCTx->saved_data["imgWidth"].toInt();

    variable_list SavedVariables = vCTx->get_saved_variables();
    torch::Tensor GaussianIdsSorted = SavedVariables[0];
    torch::Tensor TileBins = SavedVariables[1];
    torch::Tensor XYs = SavedVariables[2];
    torch::Tensor Conics = SavedVariables[3];
    torch::Tensor Colors = SavedVariables[4];
    torch::Tensor Opacity = SavedVariables[5];
    torch::Tensor Background = SavedVariables[6];
    torch::Tensor FinalTs = SavedVariables[7];
    torch::Tensor FinalIdx = SavedVariables[8];

    torch::Tensor v_outAlpha = torch::zeros_like(OutImg.index({ "...", 0 }));

    auto BackwardTensor = rasterize_backward_tensor(ImgHeight, ImgWidth,
        GaussianIdsSorted,
        TileBins,
        XYs,
        Conics,
        Colors,
        Opacity,
        Background,
        FinalTs,
        FinalIdx,
        OutImg,
        v_outAlpha);

    torch::Tensor OutXY = std::get<0>(BackwardTensor);
    torch::Tensor OutConic = std::get<1>(BackwardTensor);
    torch::Tensor OutColors = std::get<2>(BackwardTensor);
    torch::Tensor OutOpacity = std::get<3>(BackwardTensor);
    torch::Tensor None;

    return { OutXY,
            None, //Authorm comment:  depths
            None, //Authorm comment:  radii
            OutConic,
            None, //Authorm comment:  numTilesHit
            OutColors,
            OutOpacity,
            None, //Authorm comment:  imgHeight
            None, //Authorm comment:  imgWidth
            None //Authorm comment:  background
    };
}


